import torch
from model import nnModel
from cartpole_dynamic import CartPole
from argparse import ArgumentParser
from cartpole_render import CartPoleVisualizer
import pygame
import time

args = ArgumentParser()
args.add_argument("--dt", type=float, default=0.033)
args.add_argument("--polelen", type=float, default=1.5)
args.add_argument("--m_cart", type=float, default=20)
args.add_argument("--m_pole", type=float, default=10)
args.add_argument("--g", type=float, default=9.81)
args.add_argument("--timesteps", type=int, default=1000)
args.add_argument("--control_nn", type=str, default="controll_nn.pth")
args = args.parse_args()

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
state = torch.zeros((1, 4), device=device)
hx = None
dt = args.dt  # Delta t
len = args.polelen  # Length of pole
m_cart = args.m_cart  # Mass of cart
m_pole = args.m_pole  # Mass of pole
g = args.g  # Accerleration due to gravity
N = args.timesteps  # Number of time steps

model = CartPole(dt, len, m_cart, m_pole, g).to(device)
control_nn = nnModel(4, 1).to(device)
control_nn.load_state_dict(torch.load(args.control_nn))


def reset():
    global state, hx
    state = torch.zeros((1, 4), device=device)
    hx = None


visualizer = CartPoleVisualizer()
while visualizer.running:
    start_time = time.time()
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            visualizer.running = False
        if event.type == pygame.MOUSEBUTTONDOWN:
            mouse_pos = event.pos
            if visualizer.is_reset_button_clicked(mouse_pos):
                visualizer.reset_position()
                reset()
        if event.type == pygame.KEYUP:
            if event.key == pygame.K_r:
                visualizer.reset_position()
                reset()
    force, hx = control_nn(state, hx)
    state, _ = model(state, force)
    x = state[0, 0].item()
    theta = state[0, 2].item()
    visualizer.update(x, 0, 0, theta)
    end_time = time.time()
    time_cost = end_time - start_time
    pygame.time.delay(int(dt * 1000) - int(time_cost * 1000))
visualizer.quit()
